#ifndef MPQC_TENSOR_HPP
#define MPQC_TENSOR_HPP

#include <assert.h>
#include <algorithm>

#include <tiled_array.h>

#include <boost/fusion/include/vector.hpp>
#include <boost/fusion/include/push_back.hpp>
#include <boost/fusion/include/at_c.hpp>

#include "range.hpp"

#define MPQC_ENUM_PARAM_1( X ) X##1
#define MPQC_ENUM_PARAM_2( X ) MPQC_ENUM_PARAM_1( X ) , X##2
#define MPQC_ENUM_PARAM_3( X ) MPQC_ENUM_PARAM_2( X ) , X##3
#define MPQC_ENUM_PARAM_4( X ) MPQC_ENUM_PARAM_3( X ) , X##4
#define MPQC_ENUM_PARAM_5( X ) MPQC_ENUM_PARAM_4( X ) , X##5
#define MPQC_ENUM_PARAM_6( X ) MPQC_ENUM_PARAM_5( X ) , X##6
#define MPQC_ENUM_PARAM_7( X ) MPQC_ENUM_PARAM_6( X ) , X##7
#define MPQC_ENUM_PARAM_8( X ) MPQC_ENUM_PARAM_7( X ) , X##8

#define MPQC_REPEAT_1( X ) X( 1 )
#define MPQC_REPEAT_2( X ) MPQC_REPEAT_1( X ) X( 2 )
#define MPQC_REPEAT_3( X ) MPQC_REPEAT_2( X ) X( 3 )
#define MPQC_REPEAT_4( X ) MPQC_REPEAT_3( X ) X( 4 )
#define MPQC_REPEAT_5( X ) MPQC_REPEAT_4( X ) X( 5 )
#define MPQC_REPEAT_6( X ) MPQC_REPEAT_5( X ) X( 6 )
#define MPQC_REPEAT_7( X ) MPQC_REPEAT_6( X ) X( 7 )
#define MPQC_REPEAT_8( X ) MPQC_REPEAT_7( X ) X( 8 )

#define MPQC_REPEAT( N , X ) MPQC_REPEAT_##N( X )

#define MPQC_ENUM_PARAMS( N , X ) MPQC_ENUM_PARAM_##N( X )

namespace mpqc {

    template<typename T, size_t N>
    class tensor;

}


namespace mpqc {
namespace detail {
namespace tensor {

    template<class T>
    struct is_tensor : std::false_type {};

    template<typename T, size_t N>
    struct is_tensor< mpqc::tensor<T, N> > : std::true_type {};

    template<typename T, size_t N>
    struct is_tensor< const mpqc::tensor<T, N> > : std::true_type {};


    struct assign {
        template<class T, class U, class S>
        void operator()(T &t, U &u, const S &s) const {
            //std::cout << "t" << s << " = " << "u" << s << std::endl;
            t(s) = u(s);
        }
    };

    template<class T>
    class reference_wrapper {
    public:
        typedef T type;

        explicit reference_wrapper(T& t): t_(&t) { }

        operator T& () const { return *t_; }
        T& get() const { return *t_; }
        T* get_pointer() const { return t_; }

    private:
        T* t_;
    };

    template<class T>
    inline const reference_wrapper<T const> cref(T const & t) {
        return reference_wrapper<T const>(t);
    }


    template<class R, size_t K>
    struct Tie {
        struct Ref : reference_wrapper<R> {
            Ref(R &r) : reference_wrapper<R>(r) {}
        };
        Ref data[K];
        R& operator[](int i) const {
            return data[i];
        }
    };


    template<class F, class T, class U, class S>
    void apply(F f, T &t, U &u, const S &s, std::integral_constant<int,1>) {
        assert(t.dims()[0] == u.dims()[0]);
        int n = t.dims()[0];
        for (int i = 0; i < n; ++i) {
            f(t, u, boost::fusion::push_front(s, cref(i)));
        }
    }

    template<class F, class T, class U, class S, int K>
    void apply(F f, T &t, U &u, const S &s, std::integral_constant<int,K>) {
        const static int J = K-1; // col major traversal
        assert(t.dims()[J] == u.dims()[J]);
        int n = t.dims()[J];
        for (int i = 0; i < n; ++i) {
            apply(f, t, u,
                  boost::fusion::push_front(s, cref(i)),
                  std::integral_constant<int,K-1>());
        }
    }

} // namespace tensor
} // namespace detail
} // namespace mpqc



namespace mpqc {


    template<class F, class T, class U, size_t N>
    void apply(F f, tensor<T,N> t, tensor<U,N> u) {
        detail::tensor::apply(f, t, u,
                              boost::fusion::vector<>(),
                              std::integral_constant<int,N>());
    }

    enum TENSOR_STORAGE_ORDER { TENSOR_STORAGE_UNDEFINED,
                                TENSOR_COLUMN_MAJOR,
                                TENSOR_ROW_MAJOR };

    template<typename T, size_t N>
    class tensor {

    public:

        static const size_t RANK = N;

        tensor(T *data, const size_t (&dims)[N], TENSOR_STORAGE_ORDER order = TENSOR_ROW_MAJOR) :
          order_(order),
          data_(data)
        {
            for (int i = 0; i < N; ++i) {
                base_[i] = 0;
                dims_[i] = dims[i];
            }
            if (order == TENSOR_COLUMN_MAJOR) {
                size_t stride = 1;
                for (int i = 0; i < N; ++i) {
                    strides_[i] = stride;
                    stride *= dims[i];
                }
            } else {
                size_t stride = 1;
                for (int i = N - 1; i >= 0; --i) {
                    strides_[i] = stride;
                    stride *= dims[i];
                }
            }
        }

        tensor(T *data, const size_t *dims, const size_t *strides) :
          order_(TENSOR_STORAGE_UNDEFINED),
          data_(data)
        {
            for(int i = 0; i < N; ++i) {
              base_[i] = 0;
              dims_[i] = dims[i];
              strides_[i] = strides[i];
            }
        }

        template <typename U>
        void operator=(const tensor<U,N> &u) {
            apply(detail::tensor::assign(), *this, u);
        }

        const size_t* dims() const { return dims_; }

      /*  tensor<T,N> operator+(tensor<T,N> lhs, const tensor<T,N> &rhs){

        }
      */

#define MPQC_TENSOR_RANGE_OPERATOR(N)                                           \
        tensor<T,N> operator()(MPQC_ENUM_PARAMS(N, const range &r)) {           \
            detail::tensor::Tie<const range, N> t =                             \
                {{ MPQC_ENUM_PARAMS(N,r) }};                                    \
            return generate< tensor<T,N> >(*this, t);                           \
        }                                                                       \
                                                                                \
        tensor<const T,N> operator()(MPQC_ENUM_PARAMS(N, const range &r)) const { \
            detail::tensor::Tie<const range, N> t =                             \
                {{ MPQC_ENUM_PARAMS(N,r) }};                                    \
            return generate< tensor<const T,N> >(*this, t);                     \
        }                                                                       \

        MPQC_REPEAT(8, MPQC_TENSOR_RANGE_OPERATOR)



#define MPQC_TENSOR_INDEX_OPERATOR(N)                                           \
        T& operator()(MPQC_ENUM_PARAMS(N, int i)) {                             \
            detail::tensor::Tie<const int, N> t =                               \
                {{ MPQC_ENUM_PARAMS(N,i) }};                                    \
            return tensor<T,N>::operator()(t.data);                             \
        }                                                                       \
                                                                                \
        const T& operator()(MPQC_ENUM_PARAMS(N, int i)) const {                 \
            detail::tensor::Tie<const int, N> t =                               \
                {{ MPQC_ENUM_PARAMS(N,i) }};                                    \
            return tensor<T,N>::operator()(t.data);                             \
        }                                                                       \

        MPQC_REPEAT(8, MPQC_TENSOR_INDEX_OPERATOR)

        template<class Index>
        T& operator()(const Index &idx) {
            return data_[index(idx)];
        }

        template<class Index>
        const T& operator()(const Index &idx) const {
            return data_[index(idx)];
        }

    private:

        TENSOR_STORAGE_ORDER order_;
        T *data_;
        size_t base_[N];
        size_t dims_[N];
        size_t strides_[N];

        tensor() {}

        template<class Index>
        size_t index(const Index &idx, std::integral_constant<int,0>) const {
            return 0;
        }

        template<class Index, int K>
        size_t index(const Index &idx, std::integral_constant<int,K>) const {
            static const int J = K-1;
            return ((base_[J] + boost::fusion::at_c<J>(idx))*strides_[J] +
                    (index(idx, std::integral_constant<int,J>())));
        }

        template<class Index>
        size_t index(const Index &idx) const {
            assert(boost::fusion::result_of::size<Index>::value == N && "");
            size_t i = index(boost::fusion::as_vector(idx), std::integral_constant<int,N>());
            //std::cout << idx << "->" << i << std::endl;
            return i;
        }

        template<class U, size_t K>
        static U generate(U &u, detail::tensor::Tie<const range, K> &rt) {
            assert(K == RANK && "Tensor rank mismatch in range operator()");
            U t;
            t.data_ = u.data_;
            for (int i = 0; i < N; ++i) {
                const range &r = rt[i];
                t.dims_[i] = r.size();
                t.base_[i] = u.base_[i] + *r.begin();
                t.strides_[i] = u.strides_[i];
                //printf("generate: %lu %lu %lu\n", t.dims_[i], t.base_[i], t.strides_[i]);
            }
            return t;
        }

    };

}


#endif /* MPQC_TENSOR_HPP */
